import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# 设置中文字体显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

class SimpleUNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 下采样部分
        self.down1 = self._block(1, 32)
        self.down2 = self._block(32, 64)
        self.down3 = self._block(64, 128)

        # 中间层
        self.mid = nn.Sequential(
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 128, 3, padding=1),
            nn.ReLU()
        )

        # 上采样部分
        self.up1 = self._block(256, 64)
        self.up2 = self._block(128, 32)
        self.up3 = nn.Conv2d(64, 1, 3, padding=1)

    def _block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.ReLU()
        )

    def forward(self, x, t):
        x1 = self.down1(x)
        x2 = self.down2(nn.MaxPool2d(2)(x1))
        x3 = self.down3(nn.MaxPool2d(2)(x2))

        m = self.mid(nn.MaxPool2d(2)(x3))

        u1 = self.up1(torch.cat([
            nn.Upsample(scale_factor=2)(m),
            nn.Upsample(size=(6, 6))(x3)
        ], dim=1))

        u2 = self.up2(torch.cat([
            nn.Upsample(scale_factor=2)(u1),
            nn.Upsample(size=(12, 12))(x2)
        ], dim=1))

        u3 = self.up3(torch.cat([
            nn.Upsample(scale_factor=2)(u2),
            nn.Upsample(size=(24, 24))(x1)
        ], dim=1))

        return nn.Upsample(size=(28, 28))(u3)

class Diffusion:
    def __init__(self, T=1000):
        self.T = T

    def reverse_process(self, model, xt, t):
        return model(xt, t)

def load_and_generate():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 初始化模型和扩散过程
    model = SimpleUNet().to(device)
    diffusion = Diffusion()

    try:
        # 加载预训练模型
        model.load_state_dict(torch.load('best_model.pth'))
        print("模型加载成功，开始生成图像...")

        # 生成图像
        with torch.no_grad():
            sample = torch.randn(1, 1, 28, 28).to(device)
            for t in reversed(range(0, diffusion.T, 50)):  # 跳步采样加速
                t_tensor = torch.tensor([t]).to(device)
                sample = diffusion.reverse_process(model, sample, t_tensor)

            # 显示结果
            plt.figure(figsize=(8, 8))
            plt.imshow(sample.cpu().squeeze(), cmap='gray')
            plt.title("Diffusion模型生成结果")
            plt.axis('off')
            plt.show()

    except Exception as e:
        print(f"错误: {str(e)}")
        print("请确保best_model.pth文件存在且模型结构匹配")

if __name__ == "__main__":
    load_and_generate()
